{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a029ef6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d980159e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from dynamic_quant_ops import tensor_quant_softmax\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3b1f508f",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(0, 100, 1000)\n",
    "in1 = np.cos(x)\n",
    "in2 = np.cos(1.4*x)\n",
    "inputs = torch.from_numpy(np.stack((in1, in2))).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cedc3b55",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.5000, 0.5000],\n",
       "        [0.5012, 0.4988],\n",
       "        [0.5048, 0.4952],\n",
       "        ...,\n",
       "        [0.6600, 0.3400],\n",
       "        [0.7040, 0.2960],\n",
       "        [0.7427, 0.2573]], dtype=torch.float64)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.nn.functional.softmax(inputs, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4ab8b104",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.4961, 0.4961],\n",
       "        [0.5000, 0.4961],\n",
       "        [0.5039, 0.4922],\n",
       "        ...,\n",
       "        [0.6602, 0.3359],\n",
       "        [0.7031, 0.2930],\n",
       "        [0.7422, 0.2539]], dtype=torch.float64)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tensor_quant_softmax(inputs)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "trans-fat",
   "language": "python",
   "name": "trans-fat"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
